package study.算法Algorithm.赫夫曼编码;

/*
    赫夫曼编码压缩
 */

import java.io.*;
import java.util.*;

public class HuffmanCode {
    public static void main(String[] args) {
//        //测试 ：  将一个字符串进行压缩和解压
//        String str = "i like like like java do you like a java";
//        byte[] bytes = str.getBytes();
//        byte[] huffmanZip = huffmanZip(bytes);
//        System.out.println("huffmanZip = " + Arrays.toString(huffmanZip));
//        byte[] sourceByte = decode(huffmanCodes, huffmanZip);
//        System.out.println("原来的字符串 = "+new String(sourceByte));
        //!!!!!我真的sb了  上面测试代码没有注释，赫夫曼编码表一直是对应上面的，导致去解压文件时赫夫曼编码表
        //传入错误，文件一直解压不成功

//        //测试文件压缩
        zipFile("demo.jpg","c.zip");
//        //测试文件解压
        unZipFile("c.zip","demo1.jpg");
    }

    /**
     * 编写一个方法，完成对压缩文件的解压
     * @param zipFile  准备解压的文件
     * @param dstFile  将文件解压到的路径
     */
    public static void unZipFile(String zipFile,String dstFile){
        try(
            BufferedInputStream bis = new BufferedInputStream(new FileInputStream(zipFile));
            BufferedOutputStream bos = new BufferedOutputStream(new FileOutputStream(dstFile));
            //需要使用对象流去读，将数组和赫夫曼编码表读出来
            ObjectInputStream ois = new ObjectInputStream(bis)
        ) {
            //读取数据
            byte[] huffmanBytes = (byte[]) ois.readObject();
            //读取赫夫曼编码表
            Map<Byte,String> huffmanCodes = (Map<Byte, String>) ois.readObject();
            byte[] bytes = decode(huffmanCodes, huffmanBytes);
            bos.write(bytes);
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    /**
     * 编写方法，将一个文件进行压缩
     * @param srcFile  希望压缩的文件的全路径
     * @param dstFile  压缩后将压缩文件放到哪个目录下
     */
    public static void zipFile(String srcFile,String dstFile){

        try(
            BufferedInputStream bis = new BufferedInputStream(new FileInputStream(srcFile));
            BufferedOutputStream bos = new BufferedOutputStream(new FileOutputStream(dstFile));
            //需要使用对象流去写，否则压缩后的数据会和对应赫夫曼编码表混在一起
            ObjectOutputStream oos = new ObjectOutputStream(bos)
        ) {
            //创建一个和源文件大小一样的byte[]
            byte[] b = new byte[bis.available()];
            //读取文件
            bis.read(b);
            //对源文件压缩
            byte[] huffmanBytes = huffmanZip(b);
            //使用对象流将源文件压缩后的数组 写入到指定目录
            oos.writeObject(huffmanBytes);
            //一定要把 赫夫曼编码表写入到压缩文件，  否则解压时没有表，不能对应解压，数据会丢失
            oos.writeObject(huffmanCodes);
        } catch (Exception e) {
            e.printStackTrace();
        }

    }


    //将赫夫曼编码解压成原数据
    //思路：
    //1、先将赫夫曼编码数组 转换成对应的二进制字符串
    //2、将字符串 按照赫夫曼编码  转换成 原数据
    /**
     * 编写一个方法，完成对压缩数据的解码
     *
     * @param huffmanCodes 赫夫曼编码表
     * @param huffmanBytes 压缩后的赫夫曼编码数组
     * @return 原数据对应的byte数组
     */
    private static byte[] decode(Map<Byte, String> huffmanCodes, byte[] huffmanBytes) {

        //得到huffmanBytes 对应的二进制的字符串
        StringBuilder stringBuilder = new StringBuilder();
        //将byte数组转成二进制的字符串
        for (int i = 0; i < huffmanBytes.length; i++) {
            byte b = huffmanBytes[i];
            //判断是不是最后一个字节
            boolean flag = (i == huffmanBytes.length - 1);
            stringBuilder.append(byteToBitString(!flag, b));
        }
        //把字符串按照指定的赫夫曼编码表进行解码
        //把赫夫曼编码表进行调换，因为要反向查询  传进来的编码表的key，value需要调换
        Map<String, Byte> map = new HashMap<>();
        huffmanCodes.forEach((k, v) -> {
            map.put(v, k);
        });
        //创建一个集合存放byte
        List<Byte> list = new ArrayList<>();
        for (int i = 0; i < stringBuilder.length();) {
            int count = 1;  //小的计数器
            boolean flag = true;
            Byte b = null;

            while (flag) {
                //取出count个字符去比较
                String key = stringBuilder.substring(i, i + count);  //i不动，让count移动，直到匹配到一个字符
                b = map.get(key);
                if (b == null) { //说明没有匹配到
                    count++;
                } else {   //说明匹配到了
                    flag = false;
                }
            }
            list.add(b);
            i += count;  //i直接移动到count的位置
        }
        //当for循环结束后，list中就存放了所有的字符
        //把list中的数据放入到byte[] 并返回
        byte b[] = new byte[list.size()];
        for (int i = 0; i < b.length; i++) {
            b[i] = list.get(i);
        }
        return b;
    }

    /**
     * 将一个byte 转成一个二进制的字符串
     *
     * @param flag 标识是否需要补高位如果是true，表示需要补高位，如果是false表示不补.如果是最后一个字节不需要补位
     * @param b    传入的byte
     * @return 返回的是该b对应的二进制字符串，（注意是按补码返回）
     */
    private static String byteToBitString(boolean flag, byte b) {
        //使用变量保存b
        int temp = b;
        //如果是整数，还需要补高位
        if (flag) {
            temp |= 256;   //按位与
        }
        String str = Integer.toBinaryString(temp);
        if (flag) {
            return str.substring(str.length() - 8);
        } else {
            return str;
        }
    }

    /**
     * 将所有方法封装起来 ,便于调用
     *
     * @param bytes 原始的字符串对应的字节数组
     * @return 经过赫夫曼编码压缩后的字节数组
     */
    private static byte[] huffmanZip(byte[] bytes) {
        //将原始字符串对应的字节数组 转为一个存储Node对象的List集合 集合中每个元素都是一个小二叉树
        List<Node> nodes = getNodes(bytes);
        //使用存储Node对象的List集合去创建一个 赫夫曼树
        Node root = createHuffmanTree(nodes);
        //使用赫夫曼树的根节点创建一个 赫夫曼编码表 map集合
        Map<Byte, String> huffmanCodes = getCodes(root);
        //调用zip（） 传入原始字符串的字节数组 和赫夫曼编码表 进行压缩
        byte[] huffmanCodeBytes = zip(bytes, huffmanCodes);
        return huffmanCodeBytes;
    }


    //编写一个方法，将字符串对应的byte数组，通过生成的赫夫曼编码表，返回一个赫夫曼编码  压缩后的byte[]

    /**
     * @param bytes        这是原始的字符串对应的byte[]
     * @param huffmanCodes 生成的赫夫曼编码  map
     * @return 返回赫夫曼编码处理后的byte[]
     */
    private static byte[] zip(byte[] bytes, Map<Byte, String> huffmanCodes) {
        //1、利用huffmanCodes 将bytes转成 赫夫曼编码对应的字符串
        StringBuilder stringBuilder = new StringBuilder();
        //遍历bytes数组
        for (byte b : bytes) {
            stringBuilder.append(huffmanCodes.get(b));
        }
        //将赫夫曼编码字符串转成byte数组
        //统计需要return的byte数组的长度
        int len;
        if(stringBuilder.length() % 8 == 0) {
			len = stringBuilder.length() / 8;
		} else {
			len = stringBuilder.length() / 8 + 1;
		}
        //创建存储压缩后的byte数组
        byte[] huffmanCodeBytes = new byte[len];
        int index = 0;    //记录时第几个byte
        for (int i = 0; i < stringBuilder.length(); i += 8) {  //因为是8位对应一个byte，所以步长是+8
            String strByte;
            if (i + 8 > stringBuilder.length()) {  //说明最后一段不够8位
                strByte = stringBuilder.substring(i);
            } else {
                strByte = stringBuilder.substring(i, i + 8);
            }
            //将strByte转成一个byte数组，放入到huffmanCodeBytes
            huffmanCodeBytes[index] = (byte) Integer.parseInt(strByte, 2);
            index++;
        }
        return huffmanCodeBytes;
    }

    /**
     * @param bytes 接收一个字节数组
     * @return 返回的就是List形式
     */
    private static List<Node> getNodes(byte[] bytes) {

        //1、创建一个ArrayList
        ArrayList<Node> nodes = new ArrayList<>();
        //存储每个byte出现的次数  -》 map
        Map<Byte, Integer> counts = new HashMap<>();
        for (byte b : bytes) {
            Integer count = counts.get(b);
            if (count == null) {  //map还没有字符数据
                counts.put(b, 1);
            } else {
                counts.put(b, count + 1);
            }
        }

        //把每个键值对 转成一个Node对象，并加入到ArrayList集合
        counts.forEach((k, v) -> {
            nodes.add(new Node(k, v));
        });
        return nodes;
    }

    //通过List 创建对应的赫夫曼树
    private static Node createHuffmanTree(List<Node> nodes) {

        while (nodes.size() > 1) {
            //排序  从小到大
            Collections.sort(nodes);
            //取出两颗最小的二叉树
            Node leftNode = nodes.get(0);
            Node rightNode = nodes.get(1);
            //创建一颗新的二叉树   父节点是两个小二叉树之和，且父节点没有data，只有权值
            Node parent = new Node(null, leftNode.wight + rightNode.wight);
            parent.left = leftNode;
            parent.right = rightNode;

            //将处理过的两颗二叉树从 nodes中删除
            nodes.remove(leftNode);
            nodes.remove(rightNode);
            //将新的二叉树加入到nodes中
            nodes.add(parent);
        }
        //将nodes中最后剩下的节点返回  ，它就是赫夫曼树的根节点
        return nodes.get(0);
    }

    //为了调用方便，重置getCodes
    private static Map<Byte, String> getCodes(Node root) {
        if (root == null) {
            return null;
        }
        //处理左子树
        getCodes(root.left, "0", stringBuilder);
        getCodes(root.right, "1", stringBuilder);

        return huffmanCodes;
    }

    //生成赫夫曼树对应的赫夫曼编码
    //思路
    //1、将赫夫曼编码表存放在Map<Byte,String>
    static Map<Byte, String> huffmanCodes = new HashMap<>();
    //2、在生成赫夫曼编码表时，需要去拼接路径，定义一个StringBuilder  存储某个叶子节点的路径
    static StringBuilder stringBuilder = new StringBuilder();

    /**
     * 功能：得到传入的node节点的所有叶子节点的赫夫曼编码，并放入到一个集合中
     *
     * @param node 传入的节点， 根节点
     * @param code 代表此节点的路径编码   向左为0  向右为1
     */
    private static void getCodes(Node node, String code, StringBuilder stringBuilder) {
        StringBuilder stringBuilder1 = new StringBuilder(stringBuilder);
        //code加入到stringBuilder1中
        stringBuilder1.append(code);
        if (node != null) {  //如果node==null不处理
            //判断当前node时叶子节点还是非叶子节点
            if (node.data == null) {  //非叶子节点
                //向左递归
                getCodes(node.left, "0", stringBuilder1);
                //向右递归
                getCodes(node.right, "1", stringBuilder1);
            } else {
                //叶子节点  表示找到了某个叶子节点的最后
                huffmanCodes.put(node.data, stringBuilder1.toString());
            }
        }
    }

    //前序遍历
    private static void preOrder(Node root) {
        if (root != null) {
            root.preOrder();
        } else {
            System.out.println("赫夫曼树为空");
        }
    }
}

//创建Node  ，带数据和权值
class Node implements Comparable<Node> {
    Byte data;  //存放数据
    int wight;  //权值，表示字符出现的次数
    Node left;
    Node right;

    public Node(Byte data, int wight) {
        this.data = data;
        this.wight = wight;
    }

    @Override
    public int compareTo(Node o) {
        return this.wight - o.wight;
    }

    @Override
    public String toString() {
        return "Node{" +
                "data=" + data +
                ", wight=" + wight +
                '}';
    }

    //前序遍历
    public void preOrder() {
        System.out.println(this);
        if (this.left != null) {
            this.left.preOrder();
        }
        if (this.right != null) {
            this.right.preOrder();
        }
    }
}
